from langchain.llms import OpenAI
from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain,SimpleSequentialChain
from langchain.chains.router.llm_router import (LLMRouterChain,RouterOutputParser)
from langchain.chains.router.multi_prompt_prompt import MULTI_PROMPT_ROUTER_TEMPLATE


script_llm = OpenAI(
    api_key='hk-mtiquv1000041663a49a34520ad3294132cbce7abf1c2ef3',
    base_url="https://api.openai-hk.com/v1"
)

physics_prompt_tpl = PromptTemplate.from_template(
    '你是个优秀的物理学家，你很擅长用简明易懂的方式回答有关物理学的问题'
    '请用中文帮我解答下列问题：\n{input}'
)

math_prompt_tpl = PromptTemplate.from_template(
    '你是个优秀的数学家，你很擅长用简明易懂的方式回答有关数学的问题'
    '请用中文帮我解答下列问题：\n{input}'
)

#模板信息列表
prompt_infos = [
    {
        'name':'physcis',
        'description':'用户解答物理相关问题',
        'prompt_template':physics_prompt_tpl
    },
    {
        'name':'math',
        'description':'用户解答数学相关问题',
        'prompt_template':math_prompt_tpl
    }
]
destion_chains = {}
for p_info in prompt_infos:
    name = p_info['name']
    prompt = p_info['prompt_template']
    chain = LLMChain(llm = script_llm,prompt=prompt)
    destion_chains [name ] = chain

